import tensorflow as tf
import numpy as np


def input_fn(filename):
    feature_a = []
    feature_b = []
    feature_c = []
    feature_d = []
    labels = []

    f = open(filename, 'r')
    lines = f.readlines()
    for line in lines:
        line_arr = line.split(" ")
        feature_a.append(int(line_arr[0]))
        feature_b.append(int(line_arr[1]))
        feature_c.append(int(line_arr[2]))
        feature_d.append(int(line_arr[3]))
        labels.append(int(line_arr[4]))
    f.close()

    features = {'aligns': np.array(feature_a, dtype=int),
                'borders': np.array(feature_b, dtype=int),
                'cols': np.array(feature_c, dtype=int),
                'dots': np.array(feature_d, dtype=int)}
    return features, np.array(labels)


classifier = tf.estimator.DNNClassifier(
    feature_columns=[tf.feature_column.numeric_column(key="aligns"),
                     tf.feature_column.numeric_column(key="borders"),
                     tf.feature_column.numeric_column(key="cols"),
                     tf.feature_column.numeric_column(key="dots")],
    hidden_units=[10, 10],
    n_classes=2)

classifier.train(
    input_fn=lambda: input_fn("data.txt"),
    steps=100)

print("train finished")
eval_result = classifier.evaluate(
    input_fn=lambda: input_fn("test.txt"), steps=100)
print('\nTest set accuracy: {accuracy:0.3f}\n'.format(**eval_result))
